# material_scanner.py
import os
import sqlite3
import logging
import tkinter as tk
from tkinter import filedialog, simpledialog
from albailianb import generate_summary  # 假设albailian.py在同级目录

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='[material_scanner] %(asctime)s - %(levelname)s - %(message)s'
)

# 支持的媒体格式
MEDIA_EXTENSIONS = {'.mp4', '.avi', '.mov', '.mkv', '.flv', '.jpg', '.jpeg', '.png'}
DB_PATH = "material_db/material_db.sqlite"

def init_db():
    """初始化数据库和两个表结构"""
    logging.info("开始初始化数据库")
    try:
        # 创建数据库目录
        os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)

        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()

        # 创建分类表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS categories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT UNIQUE NOT NULL
            )
        """)

        # 创建媒体库表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS material_library (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                path TEXT NOT NULL,
                category_id INTEGER NOT NULL,
                summary TEXT,
                FOREIGN KEY (category_id) REFERENCES categories(id)
            )
        """)

        conn.commit()
        logging.info(f"[初始化] 数据库已创建: {DB_PATH}")

    except Exception as e:
        logging.error(f"[初始化] 数据库初始化失败: {e}")
        raise
    finally:
        if 'conn' in locals():
            conn.close()

def scan_and_save_to_db(root_path, category_name):
    """扫描媒体文件并生成摘要存储"""
    logging.info(f"开始扫描: {root_path} (分类: {category_name})")

    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()

        # 插入或获取分类ID
        cursor.execute("INSERT OR IGNORE INTO categories (name) VALUES (?)", (category_name,))
        cursor.execute("SELECT id FROM categories WHERE name = ?", (category_name,))
        category_id = cursor.fetchone()[0]

        media_count = 0

        # 递归扫描文件夹
        for dirpath, _, filenames in os.walk(root_path):
            for filename in filenames:
                if os.path.splitext(filename)[1].lower() in MEDIA_EXTENSIONS:
                    file_path = os.path.join(dirpath, filename)

                    try:
                        # 生成摘要
                        summary = generate_summary(file_path, "sk-4279edafb3314006bf64d4bb94be2e4e")

                        # 存储媒体信息
                        cursor.execute(
                            "INSERT INTO material_library (path, category_id, summary) VALUES (?, ?, ?)",
                            (file_path, category_id, summary)
                        )
                        media_count += 1
                        logging.info(f"处理成功: {file_path}")

                    except Exception as e:
                        logging.error(f"处理失败: {file_path} - {str(e)}")
                        continue

        conn.commit()
        logging.info(f"扫描完成，共处理 {media_count} 个媒体文件")

    except Exception as e:
        logging.error(f"扫描过程出错: {str(e)}")
        conn.rollback()
    finally:
        if 'conn' in locals():
            conn.close()

def select_and_build_material_db():
    """整合用户交互流程"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口

    # 1. 选择文件夹
    folder_path = filedialog.askdirectory(title="选择素材根目录")
    if not folder_path:
        logging.warning("未选择文件夹")
        return None

    # 2. 输入分类名称
    category = simpledialog.askstring("分类输入", "请输入素材分类名称:")
    if not category:
        logging.warning("未输入分类名称")
        return None

    logging.info(f"用户选择: 文件夹={folder_path}, 分类={category}")

    # 3. 初始化数据库和扫描
    init_db()
    scan_and_save_to_db(folder_path, category)
    return True

# ========================
# 可选：数据库查询示例
# ========================

def search_materials_by_category(category_name):
    """根据分类名称查询媒体文件"""
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        cursor.execute("""
            SELECT ml.path, ml.summary 
            FROM material_library ml
            JOIN categories c ON ml.category_id = c.id
            WHERE c.name = ?
        """, (category_name,))
        return cursor.fetchall()
    except Exception as e:
        logging.error(f"查询失败: {str(e)}")
        return []
    finally:
        if 'conn' in locals():
            conn.close()

def get_all_categories():
    """获取所有可用分类"""
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        cursor.execute("SELECT name FROM categories")
        return [row[0] for row in cursor.fetchall()]
    except Exception as e:
        logging.error(f"获取分类失败: {str(e)}")
        return []
    finally:
        if 'conn' in locals():
            conn.close()

# 示例用法
if __name__ == "__main__":
    # 启动图形界面
    select_and_build_material_db()

    # 示例查询
    print("可用分类:", get_all_categories())
    category = input("输入要查询的分类名称: ")
    results = search_materials_by_category(category)
    print(f"\n{category} 分类下的素材:")
    for path, summary in results:
        print(f"路径: {path}\n摘要: {summary}\n{'-'*50}")
